{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f8ca16613a0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATsAAAEvCAYAAAA6m2ZKAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAASRElEQVR4nO3cf+xddX3H8edrFf7QkaGlIJQizjQkaCYj33Q65oLzV2mIbGbZ2pjJ1KRjkWQmWzI2E+efc4sucRibOgm6OHSLomQWgRgTNBG1kAJlgFTCQv12VDEDGdtY2Xt/fE+z29tz+/323Hu/fMvn+Uhu7rnn8znnvPu537x6zr33fFJVSNKL3c+90AVI0mow7CQ1wbCT1ATDTlITDDtJTTDsJDXhJS90AX3OesW6unDTaSe93Q/ue+kcqpF0qvgv/oPn6r/T17Ymw+7CTafxvds2nfR27zjvktkXI+mU8d36xsQ2L2MlNWGqsEuyNcnDSQ4kua6nPUk+0bXfl+TSaY4nSUMNDrsk64BPAlcAFwM7klw81u0KYHP32Al8aujxJGka05zZbQEOVNWjVfUc8AXgqrE+VwGfqyV3AWcmOXeKY0rSINOE3Ubg8ZHXB7t1J9tHkuZumrDr+3p3fAqVlfRZ6pjsTLI3yd4fP/n8FGVJ0vGmCbuDwOjvQ84HFgf0AaCqdlfVQlUtbFi/boqyJOl404Td94HNSV6d5HRgO3DLWJ9bgPd038q+AXiqqg5NcUxJGmTwj4qr6kiSa4HbgHXADVX1QJJruvZdwB5gG3AAeBZ47/QlS9LJm+oOiqraw1Kgja7bNbJcwAemOYYkzYJ3UEhqgmEnqQlrciKAoW5b3HfS2zh5gNQGz+wkNcGwk9QEw05SEww7SU0w7CQ1wbCT1ATDTlITDDtJTTDsJDXBsJPUBMNOUhMMO0lNeFFNBDDEkMkDwAkEpFONZ3aSmmDYSWqCYSepCYadpCYYdpKaYNhJaoJhJ6kJhp2kJhh2kpowOOySbEryzSQPJnkgyR/19Lk8yVNJ9nWPD09XriQNM83tYkeAP66qe5KcAdyd5I6q+pexft+qqiunOI4kTW3wmV1VHaqqe7rlnwEPAhtnVZgkzdJMPrNLciHwy8B3e5rfmOTeJLcmee0sjidJJ2vqWU+S/DzwJeCDVfX0WPM9wKuq6pkk24CvAJsn7GcnsBPggo1rfzKWIbOlOFOK9MKZ6swuyWksBd3nq+rL4+1V9XRVPdMt7wFOS3JW376qandVLVTVwob166YpS5KOM823sQE+AzxYVR+f0OeVXT+SbOmO9+TQY0rSUNNcL14G/B5wf5J93bo/By4AqKpdwG8Df5jkCPCfwPaqqimOKUmDDA67qvo2kGX6XA9cP/QYkjQr3kEhqQmGnaQmGHaSmmDYSWqCYSepCYadpCYYdpKaYNhJasLav+P+RWTI5AHgBALSLHhmJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCYSepCYadpCYYdpKaYNhJaoJhJ6kJznpyCnC2FGl6ntlJaoJhJ6kJU4VdkseS3J9kX5K9Pe1J8okkB5Lcl+TSaY4nSUPN4jO7N1fVTya0XQFs7h6/Anyqe5akVTXvy9irgM/VkruAM5OcO+djStJxpg27Am5PcneSnT3tG4HHR14f7NZJ0qqa9jL2sqpaTHI2cEeSh6rqzpH29GxTfTvqwnInwAUb/UWMpNma6syuqha758PAzcCWsS4HgU0jr88HFifsa3dVLVTVwob166YpS5KOMzjskrwsyRlHl4G3A/vHut0CvKf7VvYNwFNVdWhwtZI00DTXi+cANyc5up9/qKqvJ7kGoKp2AXuAbcAB4FngvdOVK0nDDA67qnoUeH3P+l0jywV8YOgxJGlWvINCUhMMO0lN8DceL2JDZktxphS9WHlmJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCYSepCYadpCYYdpKaYNhJaoITAegYQyYPACcQ0NrnmZ2kJhh2kppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCYSepCYadpCYMDrskFyXZN/J4OskHx/pcnuSpkT4fnrpiSRpg8O1iVfUwcAlAknXAj4Cbe7p+q6quHHocSZqFWV3GvgX4YVX964z2J0kzNauw2w7cNKHtjUnuTXJrktfO6HiSdFKmnvUkyenAO4E/62m+B3hVVT2TZBvwFWDzhP3sBHYCXLDRyVhONUNmS3GmFK2mWZzZXQHcU1VPjDdU1dNV9Uy3vAc4LclZfTupqt1VtVBVCxvWr5tBWZL0/2YRdjuYcAmb5JVJ0i1v6Y735AyOKUknZarrxSQvBd4G/MHIumsAqmoX8NvAHyY5AvwnsL2qappjStIQU4VdVT0LrB9bt2tk+Xrg+mmOIUmz4B0Ukppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCYSepCYadpCZ4x71eMEMmDwAnENAwntlJaoJhJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCYSepCYadpCYYdpKaYNhJaoKznuiU42wpGsIzO0lNMOwkNWHZsEtyQ5LDSfaPrHtFkjuSPNI9v3zCtluTPJzkQJLrZlm4JJ2MlZzZ3QhsHVt3HfCNqtoMfKN7fYwk64BPAlcAFwM7klw8VbWSNNCyYVdVdwI/HVt9FfDZbvmzwG/2bLoFOFBVj1bVc8AXuu0kadUN/czunKo6BNA9n93TZyPw+Mjrg906SVp18/yCIj3ramLnZGeSvUn2/vjJ5+dYlqQWDQ27J5KcC9A9H+7pcxDYNPL6fGBx0g6randVLVTVwob16waWJUn9hobdLcDV3fLVwFd7+nwf2Jzk1UlOB7Z320nSqlvJT09uAr4DXJTkYJL3A38JvC3JI8DbutckOS/JHoCqOgJcC9wGPAj8Y1U9MJ9/hiSd2LK3i1XVjglNb+npuwhsG3m9B9gzuDpJmhHvoJDUBMNOUhOc9UTNGDJbijOlvHh4ZiepCYadpCYYdpKaYNhJaoJhJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmGHaSmmDYSWqCEwFIJzBk8gBwAoG1yDM7SU0w7CQ1wbCT1ATDTlITDDtJTTDsJDXBsJPUBMNOUhMMO0lNWDbsktyQ5HCS/SPr/jrJQ0nuS3JzkjMnbPtYkvuT7Euyd4Z1S9JJWcmZ3Y3A1rF1dwCvq6pfAn4A/NkJtn9zVV1SVQvDSpSk6S0bdlV1J/DTsXW3V9WR7uVdwPlzqE2SZmYWn9m9D7h1QlsBtye5O8nOGRxLkgaZataTJB8CjgCfn9DlsqpaTHI2cEeSh7ozxb597QR2Alyw0clYdGobMluKM6XM1+AzuyRXA1cC766q6utTVYvd82HgZmDLpP1V1e6qWqiqhQ3r1w0tS5J6DQq7JFuBPwXeWVXPTujzsiRnHF0G3g7s7+srSfO2kp+e3AR8B7goycEk7weuB85g6dJ0X5JdXd/zkuzpNj0H+HaSe4HvAV+rqq/P5V8hSctY9sOxqtrRs/ozE/ouAtu65UeB109VnSTNiHdQSGqCYSepCYadpCYYdpKaYNhJaoJhJ6kJhp2kJhh2kprgHffSGjFk8gBwAoGV8sxOUhMMO0lNMOwkNcGwk9QEw05SEww7SU0w7CQ1wbCT1ATDTlITDDtJTTDsJDXBsJPUBMNOUhOc9UQ6xTlbysp4ZiepCYadpCYsG3ZJbkhyOMn+kXUfSfKjJPu6x7YJ225N8nCSA0mum2XhknQyVnJmdyOwtWf931TVJd1jz3hjknXAJ4ErgIuBHUkunqZYSRpq2bCrqjuBnw7Y9xbgQFU9WlXPAV8ArhqwH0ma2jSf2V2b5L7uMvflPe0bgcdHXh/s1knSqhsadp8CXgNcAhwCPtbTJz3ratIOk+xMsjfJ3h8/+fzAsiSp36Cwq6onqur5qvpf4NMsXbKOOwhsGnl9PrB4gn3urqqFqlrYsH7dkLIkaaJBYZfk3JGXvwXs7+n2fWBzklcnOR3YDtwy5HiSNK1l76BIchNwOXBWkoPAXwCXJ7mEpcvSx4A/6PqeB/xdVW2rqiNJrgVuA9YBN1TVA/P4R0jScpYNu6ra0bP6MxP6LgLbRl7vAY77WYokrTbvoJDUBMNOUhOc9URq1JDZUk7lmVI8s5PUBMNOUhMMO0lNMOwkNcGwk9QEw05SEww7SU0w7CQ1wbCT1ATDTlITDDtJTTDsJDXBiQAkrdiQyQNgbUwg4JmdpCYYdpKaYNhJaoJhJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmLHsHRZIbgCuBw1X1um7dF4GLui5nAv9eVZf0bPsY8DPgeeBIVS3MpGpJOkkruV3sRuB64HNHV1TV7x5dTvIx4KkTbP/mqvrJ0AIlaRaWDbuqujPJhX1tSQL8DvAbM65LkmZq2s/s3gQ8UVWPTGgv4PYkdyfZOeWxJGmwaWc92QHcdIL2y6pqMcnZwB1JHqqqO/s6dmG4E+CCjU7GIr2YDJktZdYzpQw+s0vyEuBdwBcn9amqxe75MHAzsOUEfXdX1UJVLWxYv25oWZLUa5rL2LcCD1XVwb7GJC9LcsbRZeDtwP4pjidJgy0bdkluAr4DXJTkYJL3d03bGbuETXJekj3dy3OAbye5F/ge8LWq+vrsSpeklVvJt7E7Jqz//Z51i8C2bvlR4PVT1idJM+EdFJKaYNhJaoJhJ6kJhp2kJhh2kppg2ElqgmEnqQmGnaQmeMe9pDVpyOQBW97x7MQ2z+wkNcGwk9QEw05SEww7SU0w7CQ1wbCT1ATDTlITDDtJTTDsJDXBsJPUBMNOUhMMO0lNMOwkNSFV9ULXcJwkPwb+tafpLOAnq1xOH+s4lnUcyzqOtZp1vKqqNvQ1rMmwmyTJ3qpasA7rsA7rOFlexkpqgmEnqQmnWtjtfqEL6FjHsazjWNZxrDVRxyn1mZ0kDXWqndlJ0iBrMuySbE3ycJIDSa7raU+ST3Tt9yW5dA41bEryzSQPJnkgyR/19Lk8yVNJ9nWPD8+6ju44jyW5vzvG3p721RiPi0b+nfuSPJ3kg2N95jIeSW5IcjjJ/pF1r0hyR5JHuueXT9j2hH9LM6jjr5M81I37zUnOnLDtCd/DGdTxkSQ/Ghn7bRO2nfd4fHGkhseS7Juw7czGY8Wqak09gHXAD4FfBE4H7gUuHuuzDbgVCPAG4LtzqONc4NJu+QzgBz11XA788yqMyWPAWSdon/t49LxH/8bSb5rmPh7ArwOXAvtH1v0VcF23fB3w0SF/SzOo4+3AS7rlj/bVsZL3cAZ1fAT4kxW8b3Mdj7H2jwEfnvd4rPSxFs/stgAHqurRqnoO+AJw1Vifq4DP1ZK7gDOTnDvLIqrqUFXd0y3/DHgQ2DjLY8zQ3MdjzFuAH1ZV3w+/Z66q7gR+Orb6KuCz3fJngd/s2XQlf0tT1VFVt1fVke7lXcD5Q/c/TR0rNPfxOCpJgN8Bbhq6/1lbi2G3EXh85PVBjg+ZlfSZmSQXAr8MfLen+Y1J7k1ya5LXzqmEAm5PcneSnT3tqzoewHYm/xGvxngAnFNVh2DpPybg7J4+qz0u72PpDLvPcu/hLFzbXU7fMOGyfjXH403AE1X1yIT21RiPY6zFsEvPuvGvjFfSZyaS/DzwJeCDVfX0WPM9LF3KvR74W+Ar86gBuKyqLgWuAD6Q5NfHy+zZZl7jcTrwTuCfeppXazxWajXH5UPAEeDzE7os9x5O61PAa4BLgEMsXUIeV2bPunn9HGMHJz6rm/d4HGctht1BYNPI6/OBxQF9ppbkNJaC7vNV9eXx9qp6uqqe6Zb3AKclOWvWdVTVYvd8GLiZpcuRUasyHp0rgHuq6omeOldlPDpPHL1U754P9/RZrb+Tq4ErgXdX94HUuBW8h1Opqieq6vmq+l/g0xP2v1rj8RLgXcAXJ/WZ93j0WYth931gc5JXd2cR24FbxvrcAryn+xbyDcBTRy9pZqX7zOEzwINV9fEJfV7Z9SPJFpbG88kZ1/GyJGccXWbpA/H9Y93mPh4jJv6PvRrjMeIW4Opu+Wrgqz19VvK3NJUkW4E/Bd5ZVc9O6LOS93DaOkY/o/2tCfuf+3h03go8VFUH+xpXYzx6rea3ISt9sPTt4g9Y+uboQ926a4BruuUAn+za7wcW5lDDr7F0in8fsK97bBur41rgAZa+1boL+NU51PGL3f7v7Y71goxHd5yXshRevzCybu7jwVK4HgL+h6Wzk/cD64FvAI90z6/o+p4H7DnR39KM6zjA0udgR/9Gdo3XMek9nHEdf9+99/exFGDnvhDj0a2/8ejfxEjfuY3HSh/eQSGpCWvxMlaSZs6wk9QEw05SEww7SU0w7CQ1wbCT1ATDTlITDDtJTfg/HsOrXkqJhdMAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 360x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "import math                                   # 导入数学工具包\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def gen_mask(max_len : int) -> torch.Tensor:\n",
    "    \"\"\"_summary_\n",
    "        生成向后这样的掩码张量, 参数size是\n",
    "    Args:\n",
    "        size (int): 掩码张量的尺寸\n",
    "    \"\"\"\n",
    "    attn_shape = (1, max_len, max_len)\n",
    "    \"\"\"_summary_\n",
    "        >>> size = 5\n",
    "        >>> attn_shape = (1, max_len, max_len)\n",
    "        # k = 0 表示主对角线, k > 0 表示向右上移动\n",
    "        >>> mask       = np.triu(np.ones(shape=attn_shape), k = 1).astype(np.uint8)\n",
    "        >>> mask\n",
    "            array([[[0, 1, 1, 1, 1],\n",
    "                    [0, 0, 1, 1, 1],\n",
    "                    [0, 0, 0, 1, 1],\n",
    "                    [0, 0, 0, 0, 1],\n",
    "                    [0, 0, 0, 0, 0]]], dtype=uint8)\n",
    "        >>> torch.from_numpy(1 - mask)\n",
    "            tensor([[[1, 0, 0, 0, 0],\n",
    "                    [1, 1, 0, 0, 0],\n",
    "                    [1, 1, 1, 0, 0],\n",
    "                    [1, 1, 1, 1, 0],\n",
    "                    [1, 1, 1, 1, 1]]], dtype=torch.uint8)\n",
    "    \"\"\"\n",
    "    # TODO: 向后遮掩\n",
    "    mask   = np.triu(np.ones(shape=attn_shape), k = 1).astype(np.uint8)\n",
    "    return torch.from_numpy( 1- mask)\n",
    "\n",
    "max_len = 5\n",
    "plt.figure(figsize=(5, 5))\n",
    "plt.imshow(gen_mask(20)[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
